/*
* Copyright (C) 2020 ~ 2024 chaigec All Rights Reserved.
*
* Author:     chaigec <chaigec@163.com>
* Maintainer: chaigec <chaigec@163.com>

 * Use of this source code is governed by MIT license that can be found in the
 * LICENSE file in the root of the source tree. All contributing project authors
 * may be found in the AUTHORS file in the root of the source tree.
*/

#include "MkRtspPublisher.h"
#include "MkSdp/MkSdp.h"
#include "MkMedia/MkMediaSource.h"
#include "MkRtp/MkRtpServerPool.h"
#include "MkRtspSession.h"
#include "MkUtil/MkLog.h"

#define PusherOutChunkSize 60000
CMkRtspPublisher::CMkRtspPublisher(CMkEvent* pEvent, CMkTimer* pTimer, MkVector<CMkRtpUdpServer*>& VecRtpUdpServer, const MkEventRemovedCallback& RemoveCb)
    : CMkRtspClient(pEvent, pTimer, VecRtpUdpServer, RemoveCb)
    , m_ReadyCb(nullptr)
    , m_nPacketMaxLen(Len2K)
    , m_nPacketOffset(0)
    , m_pPacketBuffer(nullptr)
    , m_bReady(FALSE)
    , m_bSendCache(FALSE)
{
    m_pPacketBuffer = (Uint8*)malloc(m_nPacketMaxLen);
}

Uint32 CMkRtspPublisher::Start(const MkString& RtspUrl, BOOL bTcpStream, const CMkVideoParam& VideoParam, const CMkAudioParam& AudioParam, const MkRtspPusherReadyCallback& ReadyCb)
{
    m_ReadyCb = ReadyCb;
    MkString Host;
    Uint16 Port;
    MkString Path;
    MkString UserName;
    MkString Password;
    Uint32 ErrorCode = CMkRtspUtil::GetInstance().RtspUrlParse(RtspUrl, Host, Port, UserName, Password, Path);
    MkCheckErrorLog(ErrorCode, "rtsp url parse error rtspurl:%s\n", RtspUrl.c_str());
    m_VideoParam = VideoParam;
    m_AudioParam = AudioParam;
    m_bStreamTcp = bTcpStream;
    if (MkCodecUnknown != m_VideoParam.GetCodec()) {
        CMkSdpMedia SdpMedia;
        SdpMedia.SetMediaType("video");
        SdpMedia.SetProtocol("RTP/AVP");
        SdpMedia.SetPort(0);
        CMkSdpFmt SdpFmt;
        SdpFmt.SetPayload(96);
        SdpFmt.SetEncodingName(CMkMediaUtil::GetInstance().GetStringCodec(m_VideoParam.GetCodec()));
        SdpFmt.SetClockRate(MkDefaultClockRate);
        SdpFmt.SetSps(VideoParam.GetSps());
        SdpFmt.SetPps(VideoParam.GetPps());
        SdpFmt.SetVps(VideoParam.GetVps());
        SdpFmt.SetH264PacketMode(1);
        SdpFmt.SetH264ProfileLevelId(0x4D001F);
        SdpMedia.AddFmt(SdpFmt);
        SdpMedia.SetControlUri(RtspUrl + "/streamid=video");
        CMkRtspSubSession* pSubSession = new CMkRtspSubSession(Host, nullptr);
        pSubSession->SetSdpMedia(SdpMedia);
        if (!m_bStreamTcp) {
            pSubSession->SetRtpUdpServer(GetFreeRtpUdpServer());
        }
        m_MapSession[m_VideoParam.GetCodec()] = pSubSession;
    }
    if (MkCodecUnknown != m_AudioParam.GetCodec()) {
        CMkSdpMedia SdpMedia;
        SdpMedia.SetMediaType("audio");
        SdpMedia.SetProtocol("RTP/AVP");
        SdpMedia.SetPort(0);
        CMkSdpFmt SdpFmt;
        SdpFmt.SetPayload(97);
        SdpFmt.SetEncodingName(CMkMediaUtil::GetInstance().GetStringCodec(m_AudioParam.GetCodec()));
        SdpFmt.SetClockRate(CMkMediaUtil::AudioSampleRateIndexToRate(m_AudioParam.GetSampleRate()));
        SdpFmt.SetEncodingParams(m_AudioParam.GetChannel());
        SdpFmt.SetH264ProfileLevelId(1);
        SdpMedia.AddFmt(SdpFmt);
        SdpMedia.SetControlUri(RtspUrl + "/streamid=audio");
        CMkRtspSubSession* pSubSession = new CMkRtspSubSession(Host, nullptr);
        pSubSession->SetSdpMedia(SdpMedia);
        if (!m_bStreamTcp) {
            pSubSession->SetRtpUdpServer(GetFreeRtpUdpServer());
        }
        m_MapSession[m_AudioParam.GetCodec()] = pSubSession;
    }
    CMkSdp Sdp;
    Sdp.SetVersion("0");
    Sdp.SetOrigin(CMkOrigin("-", "0", "0", "IN", "IP4", "127.0.0.1"));
    Sdp.SetSessionName("mk rtsp publisher");
    Sdp.SetConnection(CMkSdpConnection("IN", "IP4", Host));
    MkUnorderedMap<Uint32, CMkRtspSubSession*>::iterator it = m_MapSession.begin();
    for (; it != m_MapSession.end(); it++) {
        Sdp.AddMedia(it->second->GetSdpMedia());
    }
    Sdp.ToString(m_StrSdp);
    return ConnectRtspUrl(RtspUrl);
}

Uint32 CMkRtspPublisher::SendEsPacket(const CMkMediaSource* pSource, const MkEsPacket& EsPacket)
{
    m_nPacketOffset = 0;
    MkUnorderedMap<Uint32, CMkRtspSubSession*>::iterator it = m_MapSession.find(EsPacket.CodecType);
    if (it == m_MapSession.end()) {
        //stream not set up;
        return NoneError;
    }
    CMkRtpPacket RtpPacket;
    RtpPacket.SetSsrc(it->second->GetSsrc());
    RtpPacket.SetPayloadType(it->second->GetPayloadType());
    if (!m_bSendCache && pSource) {
        MkVector<MkCacheEsPacket>::const_iterator itCache = pSource->GetCacheEsPacket().begin();
        for (; itCache != pSource->GetCacheEsPacket().end(); itCache++) {
            RtpPacket.SetTimeStamp(itCache->Pts * (MkDefaultClockRate / 1000));
            CMkBufferList BufferList;
            itCache->CacheBufferList.GetBufferList(BufferList);
            SendRtp(it->second, RtpPacket, EsPacket);
        }
        m_bSendCache = TRUE;
    }
    RtpPacket.SetTimeStamp(EsPacket.Pts * (MkDefaultClockRate / 1000));
    return SendRtp(it->second, RtpPacket, EsPacket);
}

Uint32 CMkRtspPublisher::Stop()
{
    return SendTearDown();
}

Uint32 CMkRtspPublisher::OnResponseError(const MkRtspStatusCode& StatusCode)
{
    if (m_ReadyCb) {
        m_ReadyCb(RtspResponseErrorBegin + StatusCode);
        m_ReadyCb = nullptr;
    }
    return CMkRtspClient::OnResponseError(StatusCode);
}

Uint32 CMkRtspPublisher::OnOptionResponse(const CMkRtspResponse& Res)
{
    if (NoneError == CMkRtspClient::OnOptionResponse(Res) && GetPlaySession().empty()) {
        SendAnnounce(m_StrSdp);
    }
    return NoneError;
}

Uint32 CMkRtspPublisher::OnAnnounceResponse(const CMkRtspResponse& Res)
{
    Uint32 ErrorCode = NoneError;
    const MkString Authenticate = Res.GetHeader(MkRtspHeaderWWWAuthenticate);
    if (MkRtspStateUnauthorized == Res.GetStatusCode()) {
        return SendAnnounce(m_StrSdp);
    }
    SetPlaySession(Res.GetHeader(MkRtspHeaderSession));
    if (!m_MapSession.empty()) {
        SendSetup(m_MapSession.begin()->second, "record");
    }
    return NoneError;
}

Uint32 CMkRtspPublisher::OnSetupResponse(const CMkRtspRequest&  RtspReq, const CMkRtspResponse& Res)
{
    Uint32 ErrorCode = CMkRtspClient::OnSetupResponse(RtspReq, Res);
    if (NoneError == ErrorCode) {
        ErrorCode = SendRecord();
    }
    return ErrorCode;
}

Uint32 CMkRtspPublisher::OnRecordResponse(const CMkRtspResponse& Res)
{
    if (m_ReadyCb) {
        m_bReady = FALSE;
        m_ReadyCb(NoneError);
        m_ReadyCb = nullptr;
    }
    return CMkRtspClient::OnRecordResponse(Res);
}

Uint32 CMkRtspPublisher::SendRtp(CMkRtspSubSession* pSubSession, CMkRtpPacket& RtpPacket, const MkEsPacket& EsPacket)
{
    switch (pSubSession->GetCodec()) {
    case MkCodecH264:
        SendH264Rtp(pSubSession, RtpPacket, EsPacket);
        break;
    case MkCodecH265:
        SendH265Rtp(pSubSession, RtpPacket, EsPacket);
        break;
    case MkCodecAac:
        SendAacRtp(pSubSession, RtpPacket, EsPacket);
    default:
        break;
    }
    return NoneError;
}

Uint32 CMkRtspPublisher::SendH264Rtp(CMkRtspSubSession* pSubSession, CMkRtpPacket& RtpPacket, const MkEsPacket& EsPacket)
{
    MkListBuffer::const_iterator it = EsPacket.BufferList.GetBufferList().begin();
    for (; it != EsPacket.BufferList.GetBufferList().end(); it++) {
        //signal packet   
        if (MkMaxRtpLen > it->GetLength()) {
            RtpPacket.SetSequenceNumber(pSubSession->GetSeqNumber());
            RtpPacket.SetMarker(TRUE);
            RtpPacket.PacketHeader(m_pPacketBuffer + m_nPacketOffset, m_nPacketMaxLen - m_nPacketOffset);
            m_nPacketOffset += RtpPacket.GetHeaderLen();
            RtpPacket.AppentPayloadBufferList(*it);
            SendRtpPacket(pSubSession, RtpPacket);
            RtpPacket.ClearPayload();
        } else {
            Uint32 nOffset = 0;
            Uint8 NaluType = ((const Uint8*)it->GetBuffer())[0];
            Uint8* pBegin = nullptr;
            nOffset++;
            while (nOffset < it->GetLength()) {
                RtpPacket.SetSequenceNumber(pSubSession->GetSeqNumber());
                RtpPacket.PacketHeader(m_pPacketBuffer + m_nPacketOffset, m_nPacketMaxLen - m_nPacketOffset);
                m_nPacketOffset += RtpPacket.GetHeaderLen();
                pBegin = m_pPacketBuffer + m_nPacketOffset;
                if (1 == nOffset) {
                    //first packet
                    RtpPacket.SetMarker(FALSE);
                    pBegin[0] = MkH264PacketFuA | (NaluType & 0xE0);
                    pBegin[1] = 0x80 | (NaluType & 0x0F);
                    RtpPacket.AppentPayloadBufferList(pBegin, 2);
                    m_nPacketOffset += 2;
                    RtpPacket.AppentPayloadBufferList((const Uint8*)it->GetBuffer() + nOffset, MkMaxRtpLen);
                    nOffset += MkMaxRtpLen;
                } else if (nOffset + MkMaxRtpLen >= it->GetLength()) {
                    //end packet
                    pBegin[0] = MkH264PacketFuA | (NaluType & 0xE0);
                    pBegin[1] = 0x40 | (NaluType & 0x0F);
                    RtpPacket.SetMarker(TRUE);
                    RtpPacket.AppentPayloadBufferList(pBegin, 2);
                    m_nPacketOffset += 2;
                    RtpPacket.AppentPayloadBufferList((const Uint8*)it->GetBuffer() + nOffset, it->GetLength() - nOffset);
                    nOffset = it->GetLength();
                } else {
                    //middle packet
                    RtpPacket.SetMarker(FALSE);
                    pBegin[0] = MkH264PacketFuA | (NaluType & 0xE0);
                    pBegin[1] = 0x00 | (NaluType & 0x0F);
                    RtpPacket.AppentPayloadBufferList(pBegin, 2);
                    m_nPacketOffset += 2;
                    RtpPacket.AppentPayloadBufferList((const Uint8*)it->GetBuffer() + nOffset, MkMaxRtpLen);
                    nOffset += MkMaxRtpLen;
                }
                SendRtpPacket(pSubSession, RtpPacket);
                RtpPacket.ClearPayload();
            }
        }
    }
    return NoneError;
}

Uint32 CMkRtspPublisher::SendH265Rtp(CMkRtspSubSession* pSubSession, CMkRtpPacket& RtpPacket, const MkEsPacket& EsPacket)
{
    MkListBuffer::const_iterator it = EsPacket.BufferList.GetBufferList().begin();
    for (; it != EsPacket.BufferList.GetBufferList().end(); it++) {
        //signal packet   
        if (MkMaxRtpLen > it->GetLength()) {
            RtpPacket.SetSequenceNumber(pSubSession->GetSeqNumber());
            RtpPacket.SetMarker(TRUE);
            RtpPacket.PacketHeader(m_pPacketBuffer + m_nPacketOffset, m_nPacketMaxLen - m_nPacketOffset);
            RtpPacket.AppentPayloadBufferList(*it);
            SendRtpPacket(pSubSession, RtpPacket);
            RtpPacket.ClearPayload();
        } else { //fu
            Uint32 nOffset = 0;
            const Uint8 *pData = (const Uint8 *)it->GetBuffer();
            Uint8 NaluType = 0;
            CMkMediaUtil::Get265NaluType(*pData, NaluType);
            Uint8 LayerId = (pData[0] & 0x01) << 5 | ((pData[1] >> 3) & 0x1F);
            Uint8 Tid = pData[1] & 0x07;
            Uint8* pBegin = nullptr;
            nOffset += 2;
            while (nOffset < it->GetLength()) {
                RtpPacket.SetSequenceNumber(pSubSession->GetSeqNumber());
                m_nPacketOffset += RtpPacket.GetHeaderLen();
                pBegin = m_pPacketBuffer + m_nPacketOffset;
                pBegin[0] = ((MkHevcPacketPaci & 0x3F) << 1) | ((LayerId >> 5) & 0x01);
                pBegin[1] = ((LayerId & 0x1F) << 3) | (Tid & 0x07);
                if (2 == nOffset) {
                    //first packet
                    RtpPacket.SetMarker(FALSE);
                    pBegin[2] = 0x80 | (NaluType & 0x3F);
                    RtpPacket.AppentPayloadBufferList(pBegin, 3);
                    m_nPacketOffset += 3;
                    RtpPacket.AppentPayloadBufferList(pData + nOffset, MkMaxRtpLen);
                    nOffset += MkMaxRtpLen;
                } else if (nOffset + MkMaxRtpLen >= it->GetLength()) {
                    //end packet
                    pBegin[2] = 0x40 | (NaluType & 0x3F);
                    RtpPacket.SetMarker(TRUE);
                    RtpPacket.AppentPayloadBufferList(pBegin, 3);
                    m_nPacketOffset += 3;
                    RtpPacket.AppentPayloadBufferList(pData + nOffset, it->GetLength() - nOffset);
                    nOffset = it->GetLength();
                } else {
                    //middle packet
                    RtpPacket.SetMarker(FALSE);
                    pBegin[2] = 0x00 | (NaluType & 0x3F);
                    RtpPacket.AppentPayloadBufferList(pBegin, 3);
                    m_nPacketOffset += 3;
                    RtpPacket.AppentPayloadBufferList(pData + nOffset, MkMaxRtpLen);
                    nOffset += MkMaxRtpLen;
                }
                RtpPacket.PacketHeader(m_pPacketBuffer + m_nPacketOffset, m_nPacketMaxLen - m_nPacketOffset);
                SendRtpPacket(pSubSession, RtpPacket);
                RtpPacket.ClearPayload();
            }
        }
    }
    return NoneError;
}

Uint32 CMkRtspPublisher::SendAacRtp(CMkRtspSubSession* pSubSession, CMkRtpPacket& RtpPacket, const MkEsPacket& EsPacket)
{
    return NoneError;
}

Uint32 CMkRtspPublisher::SendRtpPacket(CMkRtspSubSession* pSubSession, const CMkRtpPacket& RtpPacket)
{
    Uint32 ErrorCode = NoneError;
    if (m_bStreamTcp) {
        Uint32 StreamLen = RtpPacket.GetPayloadLength() + RtpPacket.GetHeaderLen();
        Uint8 *pBegin = m_pPacketBuffer + m_nPacketOffset;
        pBegin[0] = 0x24;
        pBegin[1] = 0x00;
        pBegin[2] = (StreamLen >> 8) & 0xFF;
        pBegin[3] = StreamLen & 0xFF;
        CMkBufferList SendList;
        SendList.Append(pBegin, 4);
        SendList.Append(RtpPacket.GetHeaderData(), RtpPacket.GetHeaderLen());
        SendList.Append(RtpPacket.GetPayloadBufferList());
        ErrorCode = Send(SendList);
    } else {
        ErrorCode = pSubSession->SendRtpPacket(RtpPacket);
    }
    return ErrorCode;
}

CMkRtspPublisher::~CMkRtspPublisher()
{
    MkFree(m_pPacketBuffer);
}

